package glm

import (
	"chatgpt/helper"
	"chatgpt/helper/code"
	"chatgpt/logq"
	"chatgpt/model/glm/request"
	"chatgpt/plugin/elastic"
	"chatgpt/services/dataset"
	"crypto/md5"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/slclub/go-tips"
	"mime/multipart"
	"strings"
	"time"
)

// UploadDatasetFile 文件上传
func UploadDatasetFile(c *gin.Context) {
	_, fileHeader, err := c.Request.FormFile("dataset-file")

	if err != nil {
		helper.Echo(c, int(code.DATASET_UPLOAD_NIL_ERROR), "")
		logq.InfoF("ChatGLM.DATASET.UploadDatasetFile Error: ", err)
		return
	}

	fileNameArr := strings.Split(fileHeader.Filename, ".")
	fileSuffix := fileNameArr[len(fileNameArr)-1]
	if fileSuffix == "" || fileSuffixVerify(fileSuffix) == false {
		helper.Echo(c, int(code.DATASET_FILE_SUFFIX_ERROR), "")
		logq.InfoF("ChatGLM.DATASET.UploadDatasetFile Error: ", code.CODENO.String(code.DATASET_FILE_SUFFIX_ERROR))
		return
	}

	fileToDataset(c, fileSuffix, fileHeader)
}

// 验证文件后缀
func fileSuffixVerify(fileSuffix string) (res bool) {
	res = false
	suffixConf := [2]string{"txt", "xlsx"}

	for _, value := range suffixConf {
		if fileSuffix == value {
			res = true
		}
	}

	return res
}

// HandleClearDataset 清理接口
func HandleClearDataset(ctx *gin.Context) {
	dataset := &dataset.Dataset{}
	err := dataset.ClearEs()

	if err != nil {
		helper.Echo(ctx, code.DATASET_CLEAR_ERROR, nil)
		logq.InfoF("ChatGLM.DATASET.HandleClearDataset Error: ", err)
		return
	}

	logq.InfoF("ChatGLM.DATASET.HandleClearDataset Success: ", code.CODENO.String(code.DATASET_CLEAR_SUCCESS))
	helper.Count(ctx, code.SUCCESS, nil)
}

// GetDatasetFileList 获取数据集插入结果列表
func GetDatasetFileList(c *gin.Context) {
	var req request.DatasetFileList

	page := c.Query("page")
	pageSize := c.Query("pageSize")
	if page == "" {
		page = "1"
	}

	if pageSize == "" {
		pageSize = "10"
	}

	req.Page = tips.Int(page)
	req.PageSize = tips.Int(pageSize)

	dataset := &dataset.Dataset{}
	list, total, err := dataset.SelectDatasetList(req)

	if err != nil {
		logq.InfoF("ChatGLM.DATASET.GetDatasetFileList Error: ", err)
		helper.Echo(c, code.DATASET_FILE_LIST_ERROR, nil)
		return
	}

	data := map[string]any{
		"list":     list,
		"total":    total,
		"page":     req.Page,
		"pageSize": req.PageSize,
	}

	logq.InfoF("ChatGLM.DATASET.GetDatasetFileList Success: ", code.CODENO.String(code.DATASET_FILE_LIST_SUCCESS))
	helper.Echo(c, code.SUCCESS, data)
}

// 文件转数据集，并且插入ES
func fileToDataset(c *gin.Context, fileSuffix string, fileHeader *multipart.FileHeader) {
	var qaList []*elastic.Wenda
	var err error
	ds := &dataset.Dataset{}

	switch fileSuffix {
	case "txt":
		qaList, err = ds.HandleTxt(fileHeader)
	case "xlsx":
		qaList, err = ds.HandleXlsx(fileHeader)
	}

	if err != nil {
		helper.Echo(c, int(code.DATASET_FILE_HANDLER_ERROR), "")
		logq.InfoF("ChatGLM.DATASET.fileToDataset Error: ", err)
		return
	}
	dataset.UploadLock.Lock()
	defer dataset.UploadLock.Unlock()
	idList, addErr := ds.DataToEs(qaList)
	if addErr != nil {
		helper.Echo(c, int(code.DATASET_CREATE_ERROR), "")
		logq.InfoF("ChatGLM.DATASET.fileToDataset Error: ", addErr)
		return
	}

	datasetResToEs(fileHeader, len(qaList), idList)
	logq.InfoF("ChatGLM.DATASET.fileToDataset Success: ")
	helper.Echo(c, int(code.SUCCESS), "")
}

// DelDatasetAndRes 删除数据集和数据集结果记录
func DelDatasetAndRes(c *gin.Context) {
	var req request.DeleteDatasetAndRes
	fileName := c.Query("file_name")
	fileKey := c.Query("key")
	if fileName == "" || fileKey == "" {
		helper.Echo(c, int(code.DATASET_FILE_INVALID), code.DATASET_FILE_INVALID)
		logq.InfoF("ChatGLM.DATASET.DelDatasetAndRes Error: ", code.DATASET_FILE_INVALID)
		return
	}

	if tips.String(fmt.Sprintf("%X", md5.Sum([]byte(fileName)))) != fileKey {
		helper.Echo(c, int(code.DATASET_FILE_INVALID), code.DATASET_FILE_INVALID)
		logq.InfoF("ChatGLM.DATASET.DelDatasetAndRes Error: ", code.DATASET_FILE_INVALID)
		return
	}

	req.FileName = fileName
	req.MD5 = fileKey

	dataset := &dataset.Dataset{}
	err := dataset.DeleteDatasetAndRes(req)
	if err != nil {
		helper.Echo(c, int(code.DATASET_FILE_DELETE_ERROR), code.DATASET_FILE_DELETE_ERROR)
		logq.InfoF("ChatGLM.DATASET.DelDatasetAndRes Error: ", err)
		return
	}

	logq.InfoF("ChatGLM.DATASET.DelDatasetAndRes Success: ", fileName)
	helper.Echo(c, int(code.SUCCESS), "")
}

// 数据集插入结果存入ES
func datasetResToEs(fileHeader *multipart.FileHeader, total int, idList []string) {
	timeFormat := "2006-01-02 15:04:05"
	nowTime := time.Now().Format(timeFormat)
	fileNameArr := strings.Split(fileHeader.Filename, ".")

	data := elastic.WendaFile{
		FileFullName: fileHeader.Filename,
		FileName:     strings.Join(fileNameArr[:len(fileNameArr)-1], "."),
		FileSuffix:   fileNameArr[len(fileNameArr)-1],
		Total:        total,
		IdList:       idList,
		UpdatedAt:    nowTime,
		CreatedAt:    nowTime,
		MD5:          tips.String(fmt.Sprintf("%X", md5.Sum([]byte(fileHeader.Filename)))),
	}

	dataset := &dataset.Dataset{}
	err := dataset.SaveDatasetResToEs(data)

	if err != nil {
		logq.InfoF("ChatGLM.DATASET.datasetResToEs Error: ", err)
		return
	}

	logq.InfoF("ChatGLM.DATASET.datasetResToEs Success: ", code.DATASET_FILE_SAVE_SUCCESS)
}
